"""
数据库迁移脚本
用于执行SQL迁移，更新数据库结构以支持逻辑目录功能
"""
import os
import sys
import logging
import pymysql
from dotenv import load_dotenv

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 加载环境变量
load_dotenv()

def get_database_config():
    """从环境变量获取数据库配置"""
    # 从DATABASE_URL解析配置
    database_url = os.getenv('DATABASE_URL')
    if not database_url:
        logger.error("DATABASE_URL环境变量未设置")
        sys.exit(1)
    
    try:
        # 解析DATABASE_URL格式：mysql+pymysql://user:password@host:port/dbname
        parts = database_url.replace('mysql+pymysql://', '').split('@')
        user_pass = parts[0].split(':')
        host_db = parts[1].split('/')
        
        host_port = host_db[0].split(':')
        host = host_port[0]
        port = int(host_port[1]) if len(host_port) > 1 else 3306
        database = host_db[1]
        user = user_pass[0]
        password = user_pass[1] if len(user_pass) > 1 else ''
        
        return {
            'host': host,
            'port': port,
            'user': user,
            'password': password,
            'database': database,
            'charset': 'utf8mb4'
        }
    except Exception as e:
        logger.error(f"解析DATABASE_URL失败: {e}")
        sys.exit(1)

def execute_migration(config, migration_file):
    """执行SQL迁移文件"""
    logger.info(f"正在执行迁移: {migration_file}")
    
    # 读取SQL文件
    try:
        with open(migration_file, 'r', encoding='utf-8') as f:
            sql_script = f.read()
    except Exception as e:
        logger.error(f"读取迁移文件失败: {e}")
        return False
    
    # 连接数据库
    try:
        connection = pymysql.connect(**config)
        with connection.cursor() as cursor:
            # 执行SQL脚本
            cursor.execute(sql_script)
            connection.commit()
        connection.close()
        logger.info(f"迁移 {migration_file} 执行成功")
        return True
    except Exception as e:
        logger.error(f"执行迁移失败: {e}")
        return False

def main():
    """主函数"""
    logger.info("开始数据库迁移")
    
    # 获取数据库配置
    config = get_database_config()
    
    # 执行迁移
    migrations_dir = os.path.join(os.path.dirname(__file__), 'mysql', 'migrations')
    if not os.path.exists(migrations_dir):
        logger.error(f"迁移目录 {migrations_dir} 不存在")
        sys.exit(1)
    
    migration_files = [f for f in os.listdir(migrations_dir) if f.endswith('.sql')]
    migration_files.sort()  # 按文件名排序
    
    if not migration_files:
        logger.info("没有找到迁移文件")
        sys.exit(0)
    
    success_count = 0
    for migration_file in migration_files:
        file_path = os.path.join(migrations_dir, migration_file)
        if execute_migration(config, file_path):
            success_count += 1
    
    logger.info(f"迁移完成，成功执行 {success_count}/{len(migration_files)} 个迁移")

if __name__ == "__main__":
    main() 